#include "IndexNSW.h"
#include <cstddef>
#include <vector>

float fvec_l2(const float *x, const float *y, size_t d) {
    float res = 0.f;
    for (size_t i = 0; i != d; i++) {
        res += (x[i] - y[i]) * (x[i] - y[i]);
    }
    return res;
}


void IndexNSW::train() {
    neighbors_.resize(ntotal_);
    for (idx_t i = 0; i != ntotal_; i++) {
        const float *xi = get_ith_vector(i);
        if (entry_point_ == -1) {
            entry_point_ = i;
        } else {
            neighbors_[i] = search(xi, m_);
        }
    }
}

std::vector<idx_t> IndexNSW::search(const float *xq, size_t k) {
    idx_t entry_pt = entry_point_;
    float entry_dist = fvec_l2(xq, get_ith_vector(entry_pt), d_);
 
    std::vector<bool> visited(ntotal_, false);
    CandidateQueue candidate;
    ResultQueue result;

    candidate.emplace(entry_dist, entry_pt);
    result.emplace(entry_dist, entry_pt);
    visited[entry_pt] = true;


    while (!candidate.empty()) {
        auto [d, id] = candidate.top();
        candidate.pop();
        auto &neighbors_id = get_neighbors(id);

        // if (d > result.top().first) break; //may be cause bugs
        
        // iter every neighbor
        for (auto nh: neighbors_id) {
            if (visited[nh]) continue; // if current point has been visited
            visited[nh] = true;
            float dist_nh_q = fvec_l2(xq, get_ith_vector(nh), d_);
            if (result.size() < k || result.top().first > dist_nh_q) {
                result.emplace(dist_nh_q, nh);
                candidate.emplace(dist_nh_q, nh);
                if (result.size() > k) result.pop();
            }
        }
    }

    std::vector<idx_t> res;
    while (!result.empty()) {
        res.push_back(result.top().second);
        result.pop();
    }

    return res;
}